package cz.cuni.amis.episodic.bayes.experiment;

import com.google.common.collect.Table;
import com.google.common.collect.TreeBasedTable;

import cz.cuni.amis.episodic.bayes.memories.IntervalSurpriseMemoryCreator;
import cz.cuni.amis.episodic.bayes.memories.MinimizeKLMemoryCreator;
import cz.cuni.amis.episodic.bayes.memories.PointSurpriseMemoryCreator;
import cz.cuni.amis.episodic.bayes.utils.DynamicNodeProperty;
import cz.cuni.amis.episodic.bayes.utils.MemChart;
import cz.cuni.amis.episodic.bayes.utils.MemoryChartCreator;
import cz.cuni.amis.episodic.bayes.utils.NodeSnapshot;
import cz.cuni.amis.episodic.bayes.utils.chart.GraphPaintingDevice;
import cz.cuni.amis.episodic.bayes.utils.chart.PngGraphDevice;
import cz.cuni.amis.episodic.dybanem.DBN;
import cz.cuni.amis.episodic.dybanem.GoogleUtils;
import cz.cuni.amis.episodic.dybanem.RetrospectiveMemoriesCreator;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.apache.commons.lang.ArrayUtils;

import smile.Network;
import smile.learning.DataMatch;
import smile.learning.DataSet;
import smile.learning.EM;

/**
 * Memory creation in model with one layer.
 * @author ik
 */
public class Experiment_2 {
    
    public static void main(String[] args) throws FileNotFoundException, IOException, Exception {
        Network origNet = new Network();
        origNet.readFile("../datasets/1/em_01_learned_relevance.xdsl");

        List<String> evidence = Arrays.asList("A", "A", "B", "A", "A", "C", "A", "A", "C");

        DBN net = new DBN(origNet, evidence.size());

        String[] targets = new String[]{"G", "H"};
        List<DynamicNodeProperty[]> res = evolveOverTime(net, "O", evidence, targets);

        // present results
        Table<Integer, String, Double> result = format(res, targets, net);
        Double[] finalGoalEntropy = ArrayUtils.toObject(res.get(res.size() - 1)[0].getEntropy());
        result.putAll(GoogleUtils.toTable(finalGoalEntropy, "EF_G"));
        GoogleUtils.printTable(result, System.out);
        
        // draw default graphs   
        GraphPaintingDevice graphDevice = new PngGraphDevice("target/experiments/2/graphs/");
        MemoryChartCreator.createDefaultCharts(res, net, graphDevice);
        
        RetrospectiveMemoriesCreator.createMemories(net, evidence, targets, 3, new MinimizeKLMemoryCreator(graphDevice), graphDevice);
        
        learnEM(origNet, "../datasets/1/em_01_training_data.txt");
        
        //memExp(net, evidence, targets);
    }
    
    
    public static void memExp(DBN dbn, List<String> evidence, String[] targets) throws IOException {
        dbn.getNet().clearAllEvidence();
        dbn.getNet().updateBeliefs();
        NodeSnapshot[] preSnapshots = dbn.getNodesSnapshot(targets);
        dbn.setDBNEvidence("O", evidence);
        dbn.getNet().updateBeliefs();
        NodeSnapshot[] postSnapshots = dbn.getNodesSnapshot(targets);
        
       // compute KL
        double[][] kl = new double[targets.length][];
        for(int i = 0; i < kl.length; i++) {
            kl[i] = preSnapshots[i].computeKLDivergenceTo(postSnapshots[i]);
        }
        
        // draw graph
        //double[][] transpose = MemChart.transpose(kl);
        MemChart.chartToPNG(MemChart.chartXYBlockValueTime(kl, targets), "nullMem.png");
    }
    
    

    public static double[] surprise(DBN net, List<String> evidence, String target) {
        double[] surprise = new double[evidence.size()];
        for (int i = 0; i < evidence.size(); i++) {
            surprise[i] = net.surprise("O", i, evidence.get(i), target, i);
        }
        return surprise;
    }

    public static Table<Integer, String, Double> format(List<DynamicNodeProperty[]> list, String[] nodes, DBN net) {
        Table<Integer, String, Double> table = TreeBasedTable.create();
        int i = 0;
        for (DynamicNodeProperty[] props : list) {
            for (int j = 0; j < nodes.length; j++) {
                String[] outcomes = net.getNet().getOutcomeIds(nodes[j]);
                // set probabilities
                for (int k = 0; k < outcomes.length; k++) {
                    table.put(i, "P_" + nodes[j] + "_" + outcomes[k], props[j].getPosterior()[k]);
                }
                // set entropy
                double[] entropy = props[j].getEntropy();
                table.put(i, "E_" + nodes[j], entropy[entropy.length - 1]);
                // set surprise
                table.put(i, "S_" + nodes[j], props[j].getSurprise()[i]);
            }
            i++;
        }
        return table;
    }

    public static List<DynamicNodeProperty[]> evolveOverTime(DBN net, String evidenceNode, List<String> evidence, String[] targets) {
        //Table<Integer, String, DynamicNodeProperty>
       net.setDBNTargetNodes(targets);
    	List<DynamicNodeProperty[]> l = new ArrayList<>();
        for (int i = 0; i < evidence.size(); i++) {
            DynamicNodeProperty[] props = net.oneStep(evidenceNode, i, evidence.get(i), targets, i);
            l.add(props);
        }
        net.getNet().clearAllTargets();
        return l;
    }
    
    public static void learnEM(Network net, String smileDataFile) {
    	DataSet ds = new DataSet();
    	ds.readFile(smileDataFile);
		
		DataMatch[] match = ds.matchNetwork(net); //-> do i need this line? what for? It appears me as void, not creating a DataMatch[]...
		final EM em = new EM();
		em.setEqSampleSize(2); //-> learning algorithm parameter, if i'm wright
		em.setRandomizeParameters(true); //-> Not sure about this, it means that randomize the values of the nodes that are not informed in the txt?
		em.learn(ds, net, match);

    }
}
